# -*- utf-8 -*-


class Vertex:
    def __init__(self, key):
        self.id = key
        self.connect_to = {}

    def add_neighbor(self, nbr, weight=0):
        self.connect_to[nbr] = weight

    def __str__(self):
        return str(self.id) + ' connect to: ' + str([x.id for x in self.connect_to])

    def get_connections(self):
        return self.connect_to.keys()

    def get_id(self):
        return self.id

    def get_weight(self, nbr):
        return self.connect_to[nbr]


class Graph:
    def __init__(self):
        self.vertex_list = {}
        self.vertex_num = 0

    def add_vertex(self, key):
        self.vertex_num += 1
        new_vertex = Vertex(key)
        self.vertex_list[key] = new_vertex

        return new_vertex

    def get_vertex(self, key):
        if key in self.vertex_list:
            return self.vertex_list[key]
        else:
            return None

    def __contains__(self, item):
        return item in self.vertex_list

    def add_edge(self, from_key, to_key, weight=0):
        if from_key not in self.vertex_list:
            self.add_vertex(from_key)

        if to_key not in self.vertex_list:
            self.add_vertex(to_key)

        self.vertex_list[from_key].add_neighbor(self.vertex_list[to_key], weight)

    def get_vertices(self):
        return self.vertex_list.keys()

    def __iter__(self):
        return iter(self.vertex_list.values())


if __name__ == "__main__":
    g = Graph()
    for i in range(6):
        g.add_vertex(i)
    print(g.vertex_list)

    g.add_edge(0, 1, 5)
    g.add_edge(0, 5, 2)
    g.add_edge(1, 2, 4)
    g.add_edge(2, 3, 9)
    g.add_edge(3, 4, 7)
    g.add_edge(3, 5, 3)
    g.add_edge(4, 0, 1)
    g.add_edge(5, 4, 8)
    g.add_edge(5, 2, 1)

    for v in g:
        for w in v.get_connections():
            print("( %s , %s )" % (v.get_id(), w.get_id()))
